import pydicom
import pydicom.uid
import os
import nibabel as nib
import matplotlib.pyplot as plt


def read_save_dcm(dcm_folder, save_path):
    for root, dirs, files in os.walk(dcm_folder):
        idx = 0
        for f in files:
            dcm_path = os.path.join(root, f)

            dcm = pydicom.read_file(dcm_path, force=True)  # 读取dicom图像的Tags
            dcm.file_meta.TransferSyntaxUID = pydicom.uid.ImplicitVRLittleEndian
            dcm_arr = dcm.pixel_array
            plt.figure(figsize=(8, 8))  # 512 x 512
            plt.imshow(dcm_arr, 'gray')
            plt.xticks([])
            plt.yticks([])
            plt.axis('off')
            plt.savefig(os.path.join(save_path, '{}.png'.format(idx)), dpi=64)
            plt.close()
            idx += 1


def read_nii(nii_path):
    nii_img = nib.load(nii_path)
    return nii_img.get_fdata()

def save_nii(nii_path, save_path):
    nii_arr = read_nii(nii_path)
    width, height, queue = nii_arr.shape
    for i in range(queue):
        slice = nii_arr[:, :, i].transpose((1, 0))[::-1, ::-1]    # 转换nii方位和dcm一致
        # imageio.imwrite(os.path.join(save_path, '{}.png'.format(i)), slice)
        plt.figure(figsize=(8, 8))
        plt.imshow(slice, 'gray')
        plt.xticks([])
        plt.yticks([])
        plt.axis('off')
        plt.savefig(os.path.join(save_path, '{}.png'.format(i)), dpi=64)
        plt.close()


if __name__ == '__main__':

    dcm_folder_9614225_1 = 'To_SJC_gastric_data/9614225/dicom/20090917/372025/1.2.840.113619.2.55.3.12588800.5666.1253143047.577'
    dcm_folder_9614225_2 = 'To_SJC_gastric_data/9614225/dicom/20091109/383682/1.2.840.113619.2.55.3.12588800.5808.1257720980.307'
    dcm_folder_9625145_1 = 'To_SJC_gastric_data/9625145/dicom/20090825/365972/1.2.840.113619.2.55.3.12588800.5720.1251155538.300'
    dcm_folder_9625145_2 = 'To_SJC_gastric_data/9625145/dicom/20091010/375532/1.3.12.2.1107.5.1.1.20862.20091010085028893.3'
    # dcm_folder_9625145_2 = 'To_SJC_gastric_data/9625145/dicom/20091010/375532/1.3.12.2.1107.5.1.1.20862.20091010085210653.3'

    ## dcm转图片保存
    dcm_save_folder = './dcm_folder/'
    dcm_folder = [dcm_folder_9614225_1, dcm_folder_9614225_2, dcm_folder_9625145_1, dcm_folder_9625145_2]   # dcm文件路径
    dcm_save_path = ['9614225_1/', '9614225_2/', '9625145_1/', '9625145_2/']    # 保存具体分类文件夹
    if not os.path.exists(dcm_save_folder):     # 创建文件夹
        os.mkdir(dcm_save_folder)

    idx = 0  # 文件夹idx
    for folder in dcm_folder:
        save_path = os.path.join(dcm_save_folder, dcm_save_path[idx])
        if not os.path.exists(save_path):  # 创建文件夹
            os.mkdir(save_path)
        read_save_dcm(folder, save_path)
        idx += 1


    ## nii转图片保存
    nii_save_folder = './nii_folder/'
    nii_folder_9614255 = 'To_SJC_gastric_data/9614225/label/'
    nii_folder_9625145 = 'To_SJC_gastric_data/9625145/label/'
    nii_folder = [nii_folder_9614255, nii_folder_9625145]
    nii_save_path = ['9614225_1/', '9614225_2/', '9625145_1/', '9625145_2/']  # 1=治疗前 2=治疗后
    if not os.path.exists(nii_save_folder):     # 创建nii总文件夹
        os.mkdir(nii_save_folder)

    idx = 0  # 文件夹idx
    for folder in nii_folder:  # 逐个遍历nii文件夹文件
        # 遍历nii文件
        for root, ds, fs in os.walk(folder):
            for f in fs:
                save_path = os.path.join(nii_save_folder, nii_save_path[idx])  # nii分文件夹
                if not os.path.exists(save_path):
                    os.mkdir(save_path)
                save_nii(os.path.join(folder, f), save_path)
                idx += 1





